#!/usr/bin/env python

import ssummolib as slib
import cPickle as pickle
import colours
import threading, Queue
import multiprocessing
import random
#import numpy as np
import sys, os.path
try:
    import matplotlib.pyplot as plt
    matplot = True
except ImportError:
    matplot = False
    sys.stderr.write("matplotlib import error. Graph-plotting functionality will not work. ")

class LocalOptions( slib.Options ):
    def __init__(self,args=None):
        slib.Options.__init__(self)
        self.singleargs = ['--draw','-replicates','-increment']
        self.multiargs = ['-in','-out']
        self.helpTxt = { '-in' : 'Input files.',
                         '-out': 'Output files. Default: standard out',
                         '-replicates': 'Number of replicates for each dataset. [5]',
                         '-increment' : 'For rarefaction curves, only sample every `-increment` species. [100]',
                         '--draw':'If results are already saved, just draw the results saved in file `-in`',
                         '-collapse-at-rank': 'Collapse results at specific rank.',
                         }
        self.options.update( { '-replicates' : 5    ,
                               '-increment'  : 1000 ,
                               '--draw'      : False,
                               '-collapse-at-rank': None,
                             } )
        self.switches = {'--draw' : False }
        if args is not None:
            self.parseArgs(args)

def get_accession_places( node, accessions={} , path='/' ):
    for key in node.keys():
        if key == 'accessions':
            for accession in node[key]:
                accessions.update( { accession : path } )
        else:
            accessions = get_accession_places( node[key] , accessions=accessions , path=os.path.join(path,key) )
    return accessions

class PlotRarefaction( ):
    def __init__(self,datasets,replicates=None,increment=None , groups=None):
        if matplot:
            self.fig = plt.figure()
            self.axis = plt.subplot( 111 )
            plt.ioff()  # Turn off interactive mode (default anyway).
            plt.xlabel( 'Number of randomly chosen sequences' )
            plt.ylabel( 'Number of distinct taxa' )
        if replicates is None:
            self.replicates = 5
        else:
            self.replicates = replicates
        if increment is None:
            self.increment = 1000
        else:
            self.increment = increment

        self.col_ind = 0 # For assigning colours.
        number_datasets = len( datasets )
        #self.datasets = datasets
        self.colours = [colours.HSL_to_HEX(colour) for colour in colours.Heat( N=number_datasets , groups=groups ) ]
        self.n = number_datasets
        self.dataset_names = []


    def plot_scatter( self, xvals , yvals, label=None ,*args , **kwargs):
        """xvals is directly given to plt.scatter(), so must be compatible with 
        matplotlib.
        y_vals can optionally be a dictionary of vectors, in which case all the
        keys must be present in xvals.
        label can optionally be used to indicate the label in the legend.
        """
        all_xvals = [] ; all_yvals = []
        if type( yvals ) is dict:
            for n in xrange( len( yvals[xvals[0]] ) ):
                all_xvals.append( xvals )
                all_yvals.append( [ yvals[x][n] for x in xvals ] )
        else:
            all_xvals = xvals ; all_yvals = yvals
        collection = plt.scatter( plt.np.array(all_xvals) , plt.np.array(all_yvals) ,marker='x',color=self.colours[self.col_ind], label=label, *args, **kwargs )
        self.col_ind += 1
        return collection

    def plot_box( self, xvals , yvals , label=None , widths=None , *args, **kwargs ):
        """Given one or more x-value, will plot all given y-values
        as a box-and-whisker plot.
        This is designed to be called once per dataset, so that each dataset
        has the same colour.
        See http://matplotlib.sourceforge.net/api/pyplot_api.html#matplotlib.pyplot.boxplot
        for more information"""
        all_xvals = [] ; all_yvals = []
        if widths is None:
            widths = self.increment * 0.3
        mappings = plt.boxplot( yvals , positions=xvals, widths=widths,*args , **kwargs )
        for mapping in mappings:
            plt.setp( mappings[mapping] , color = self.colours[self.col_ind] )
        self.col_ind += 1
        return mappings

    def draw( self , x_max , y_max ):
        """Just provide the x-axis and y-axis limits. The x-axis is rounded
        up to the nearest 100, and the y-axis is rounded up to the nearest
        500.
        """
        plt.axis( [0 ,500*(1+( x_max //500)) , 0 , 100*(1+(y_max//100)) ])
        ticker = plt.matplotlib.ticker.MultipleLocator( base=self.increment * 10)
        self.axis.xaxis.set_major_locator( ticker )
        subticker = plt.matplotlib.ticker.MultipleLocator( base= self.increment * 5 )
        self.axis.xaxis.set_minor_locator( subticker )
        plt.grid( True )
        plt.show()

    def resume( self , results , order=None ,loc=4 ):
        artists = []
        data_names = []
        max_x = 0
        max_y = 0
        if order is None:
            iterator = results.iteritems()
        else:
            iterator = [ (key,results[key]) for key in order ]
        for (data_file,coords) in iterator:
            print data_file, len( coords[0] ) , len(coords[1] )
            xvals,yvals = coords
            max_x = max( max(xvals) , max_x )
            max_y = max( max(yvals) , max_y )
            data_names.append( data_file )
            mappings = self.plot_box( xvals , yvals )
            artists.append( mappings['boxes'][0] )
        plt.legend( artists , data_names , loc=loc , borderaxespad=0.2 )
        print 'max_x' ,max_x
        print 'max_y' ,max_y
        self.draw( max_x , max(max_y ))
        pass


class Loader( multiprocessing.Process ):
    def __init__( self ,results_dict,queue=None,semaphore=None,replicates=None,increment=None):
        multiprocessing.Process.__init__(self)
        self.accession_places = get_accession_places( results_dict , accessions={} )
        self.nOTUs = len( set( self.accession_places.values() ) )
        self.nSeqs = len( self.accession_places )   # X-axis
        self.results_dict = results_dict
        self.queue = queue
        self.semaphore = semaphore
        if replicates is None:
            self.replicates = 5
        else:
            self.replicates = replicates
        if increment is None:
            self.increment = 100
        else:
            self.increment = increment

    def run(self):
        """Loads the results in results_dict and populates a dictionary,
        with x-values as the keys, and y-values as the values.
        This allows multiple y-values to be stored per x-value.
        Optional queue will be used to send the results instead of
        returning them
        """
        yvals = {} 
        for number, nodes in self.rarefy( ):
            n_found = len( nodes ) 
            if number not in yvals:
                yvals.update( { number : [n_found] } )
            else:
                yvals[number].append( n_found )
        if self.queue is not None:
            self.queue.put( yvals )
        if self.semaphore is not None:
            self.semaphore.release()
        return yvals 

    def rarefy( self ):
        """Given a results dictionary, this will sample random accessions
        from that dictionary and will yield the number of accessions
        sampled as well as the number of genera / species / OTUs that
        that a random sample of that number of accessions finds.
        """
        node = self.results_dict
        accessions = self.accession_places.keys()
        n = len( accessions )
        i=1
        multipled = True
        random_queue = Queue.Queue()
        random_generator = RandomNumberThread( random_queue , n , self.replicates, self.increment )
        random_generator.start()
        
        while i < n:
            for replicate in xrange( self.replicates ):
                running_accessions = accessions[:]
                acc_places = self.accession_places.copy()
                n_left = n
                found_nodes = {}
                for ii in xrange( i ):
                    accession = running_accessions.pop( random_queue.get() )
                    found_node = acc_places.pop( accession )
                    if found_node in found_nodes:
                        found_nodes[found_node].append( accession )
                    else:
                        found_nodes.update( { found_node : [accession] } )
                    n_left -= 1   # Every time we choose a random accession, the list of accessions is 1 shorter.
                    random_queue.task_done()
                yield i , found_nodes
            i += self.increment
            continue
        random_generator.join()

class RandomNumberThread( threading.Thread ):
    def __init__(self,Q,N,replicates,increment):
        """Need a Queue.Queue object, number of accessions
        which we're rarefying, the number of replicates for
        each point, and the increment between points.
        """
        threading.Thread.__init__(self)
        self.N = N
        self.Q = Q
        self.increment = increment
        self.replicates = replicates
    def run( self ):
        i = 1
        while i < self.N:
            for replicate in xrange( self.replicates ):
                n_left = self.N
                for ii in xrange( i ):
                    rand = random.randint( 0, n_left-1 )
                    self.Q.put_nowait( rand )
                    n_left -= 1
            i += self.increment
        self.Q.join()
        return


def main(myoptions):
    artists = []
    if not matplot:
        save = {}
    if len( myoptions.regopts ) > 0:
        groups = [ len( in_files ) for in_files in [ myoptions.regopts[key] for key in sorted(myoptions.regopts.keys()) ]]
    else:
        groups = None
    if myoptions['--draw']:
        all_results = {}
        for saved_file in myoptions['-in']:
            with file( saved_file ,'rb' ) as infile:
                results = pickle.load( infile )
            all_results.update( results  )
        dirs = ['-group1', 'Sample_Data\\DDBJ\\*.pkl','-group2','Sample_Data\\EBI\\*BF.pkl','-group3','Sample_Data\\EBI\\*EU.pkl','-group4','Sample_Data\\HMP\\Stool\\*.pkl']
        new_options = LocalOptions( dirs )
        print 'new_options\n',new_options
        #new_options.parseArgs(dirs)
        #groups = [ len( in_files ) for in_files in [ new_options.regopts[key] for key in sorted(new_options.regopts.keys()) ]]
        groups = [13, 14,15,24]
        old_results = new_options.options['-in']
        print old_results
        #old_results += [ group for group in [ new_options.regopts[key] for key in sorted( new_options.regopts.keys()) ]]

        #plotter = PlotRarefaction( all_results.keys() ,replicates=myoptions['-replicates'],increment=myoptions['-increment'], groups=groups)
        plotter = PlotRarefaction( range(66) ,replicates=myoptions['-replicates'],increment=myoptions['-increment'], groups=groups)
        old_results = [ key[ key.rfind( os.path.sep )+1 : key.rfind( '.' ) ] for key in old_results ]
        plotter.resume( all_results , order = old_results )
        return
    else:
        pass

    plotter = PlotRarefaction( myoptions['-in'] , replicates=myoptions['-replicates'] , increment=myoptions['-increment'] , groups=groups )
    separate_processes = {}
    semaphore = multiprocessing.Semaphore( multiprocessing.cpu_count())
    for infile in myoptions['-in']:
        with file( infile,'rb' ) as in_handle:
            data_dict = pickle.load( in_handle )
        data_name = infile.rsplit('.',1)[0]
        if os.path.sep in data_name:
            data_name = data_name.rsplit( os.path.sep , 1 )[-1]
        print 'Loading',data_name
        separate_queue = multiprocessing.Queue()
        semaphore.acquire()
        plotter.dataset_names.append( data_name )
        separate_process = Loader( data_dict, queue=separate_queue , semaphore = semaphore, replicates=myoptions['-replicates'], increment=myoptions['-increment'])
        separate_process.start()
        separate_processes.update( { data_name : (separate_process, separate_queue,) } )

    n_done = 0
    while n_done < len( separate_processes ):
        for data_name in separate_processes:
            process, queue = separate_processes[data_name]
            try:
                data = queue.get( timeout=3 )
                print 'Got {0}'.format( data_name )
                n_done += 1
            except Queue.Empty:
                continue
            process.join()
            #values = plotter.load_results( data_dict , data_name )
            xvals = []; yvals = []
            for xval in data:
                xvals.append( xval ) ; yvals.append( data[xval] )
            if matplot:
                mappings = plotter.plot_box(xvals, yvals)
                artists.append( mappings['boxes'][0] )
            else:
                save.update( { data_name : [xvals , yvals ] } )
    if matplot:
        plt.legend( artists , myoptions['-in'] , loc=4, borderaxespad=0.2 )
        max_x = max( [load[0].nSeqs for load in separate_processes.values() ])
        max_y = max( [load[0].nOTUs for load in separate_processes.values() ])
        plotter.draw( max_x , max_y )
    if myoptions['-out'] is None:
        myoptions['-out'] = get_save_name()
    if myoptions['-out'] is not None:
        with file( myoptions['-out'][0] , 'wb' ) as outfile:
            pickle.dump( save , outfile , -1 )

def get_save_name():
    save_name = raw_input( 'Please enter a name to save the results, which can be drawn later:-\n')
    if save_name == ( 'n' or 'no' ):
        return None
    elif os.path.exists( save_name ):
        print 'That already exists. Trying again... (`no` or `n` to skip saving)'
        return get_save_name()

if __name__ == '__main__':
    args = sys.argv[1:]
    myoptions = LocalOptions(args)
    main(myoptions)
else:
    myoptions = LocalOptions(  )
